#include "../DWMRICore/Fiber.h"
#include "../DWMRICore/Scalar.h"
#include "../DWMRICore/DTI.h"

#include <teem\ell.h>

#include <omp.h>

char *fiber_pathname = NULL;
char *dst_pathname = NULL;

char **confPathname = NULL;
char **vertPathname = NULL;
int w, h, d, scale;
float stepSize, maxLength;

int kernel_size = 0;
float fiber_length = 0.0f;

void PrintUsage()
{
	printf("Usage: ComputeFSR_ST -fiber <configure_file> -dst <dst_file> -length <fiber_length> -kernel <kernel_size>\n");
}

int ParseArguments(const int argc, char *argv[])
{
	for (int i = 1; i < argc; ++i) {
		if (strcmp(argv[i], "-fiber") == 0) {
			fiber_pathname = argv[++i];
		} else if (strcmp(argv[i], "-dst") == 0) {
			dst_pathname = argv[++i];
		} else if (strcmp(argv[i], "-length") == 0) {
			fiber_length = (float)atof(argv[++i]);
		} else if (strcmp(argv[i], "-kernel") == 0) {
			kernel_size = atoi(argv[++i]);
		} else {
			return -1;
		}
	}

	if (fiber_pathname == NULL || dst_pathname == NULL || fiber_length < 0.0001f || kernel_size == 0)
		return -2;

	return 0;
}

float ApproximateFTLEST(CPoint3F *pos, const float delta, float *st)
{
	float d = 2.0f * delta;

	float matJ[3][3];
	matJ[0][0] = (pos[0].m_x - pos[1].m_x)  * d;
	matJ[1][0] = (pos[0].m_y - pos[1].m_y)  * d;
	matJ[2][0] = (pos[0].m_z - pos[1].m_z)  * d;
	matJ[0][1] = (pos[3].m_x - pos[6].m_x)  * d;
	matJ[1][1] = (pos[3].m_y - pos[6].m_y)  * d;
	matJ[2][1] = (pos[3].m_z - pos[6].m_z)  * d;
	matJ[0][2] = (pos[8].m_x - pos[17].m_x) * d;
	matJ[1][2] = (pos[8].m_y - pos[17].m_y) * d;
	matJ[2][2] = (pos[8].m_z - pos[17].m_z) * d;

	float matJT[3][3];
	matJT[0][0] = matJ[0][0];
	matJT[1][0] = matJ[0][1];
	matJT[2][0] = matJ[0][2];
	matJT[0][1] = matJ[1][0];
	matJT[1][1] = matJ[1][1];
	matJT[2][1] = matJ[1][2];
	matJT[0][2] = matJ[2][0];
	matJT[1][2] = matJ[2][1];
	matJT[2][2] = matJ[2][2];

	double matA[9];
	matA[0] = matJT[0][0] * matJ[0][0] + matJT[0][1] * matJ[1][0] + matJT[0][2] * matJ[2][0];
	matA[1] = matJT[0][0] * matJ[0][1] + matJT[0][1] * matJ[1][1] + matJT[0][2] * matJ[2][1];
	matA[2] = matJT[0][0] * matJ[0][2] + matJT[0][1] * matJ[1][2] + matJT[0][2] * matJ[2][2];

	matA[3] = matJT[1][0] * matJ[0][0] + matJT[1][1] * matJ[1][0] + matJT[1][2] * matJ[2][0];
	matA[4] = matJT[1][0] * matJ[0][1] + matJT[1][1] * matJ[1][1] + matJT[1][2] * matJ[2][1];
	matA[5] = matJT[1][0] * matJ[0][2] + matJT[1][1] * matJ[1][2] + matJT[1][2] * matJ[2][2];
	
	matA[6] = matJT[2][0] * matJ[0][0] + matJT[2][1] * matJ[1][0] + matJT[2][2] * matJ[2][0];
	matA[7] = matJT[2][0] * matJ[0][1] + matJT[2][1] * matJ[1][1] + matJT[2][2] * matJ[2][1];
	matA[8] = matJT[2][0] * matJ[0][2] + matJT[2][1] * matJ[1][2] + matJT[2][2] * matJ[2][2];

	st[0] = (float)matA[0];
	st[1] = (float)matA[1];
	st[2] = (float)matA[2];
	st[3] = (float)matA[4];
	st[4] = (float)matA[5];
	st[5] = (float)matA[8];

	double eigenValues[3];
	ell_3m_eigenvalues_d(eigenValues, matA, AIR_TRUE);
	double ftle = sqrt(eigenValues[0]);
	return (float)ftle;
}


void ComputeFSR_Tensor(CFiber **fibers, const uint pos, const float delta, float *st, CFiber *base)
{
	CPoint3F forwardPos[26];
	CPoint3F backwardPos[26];

	CVector3F dir = base->m_seed.dir;

	for (int i = 0; i < 26; ++i) {
		CVector3F temp_dir = fibers[i+1]->m_seed.dir;
		if (InnerProduct(dir, temp_dir) > 0.0f ) {
			int index = (fibers[i+1]->m_fCount-1) > pos ? pos : (fibers[i+1]->m_fCount-1);
			forwardPos[i] = fibers[i+1]->m_pF[index];

			index = (fibers[i+1]->m_bCount-1) > pos ? pos : (fibers[i+1]->m_bCount-1);
			backwardPos[i] = fibers[i+1]->m_pB[index];
		} else {
			int index = (fibers[i+1]->m_fCount-1) > pos ? pos : (fibers[i+1]->m_fCount-1);
			backwardPos[i] = fibers[i+1]->m_pF[index];

			index = (fibers[i+1]->m_bCount-1) > pos ? pos : (fibers[i+1]->m_bCount-1);
			forwardPos[i] = fibers[i+1]->m_pB[index];
		}
	}

	float forwardST[6]; float backwardST[6];
	float f_ftle = ApproximateFTLEST(forwardPos, delta, forwardST);
	float b_ftle = ApproximateFTLEST(backwardPos, delta, backwardST);

	st[0] = 1.0f;
	if (f_ftle > b_ftle) {
		st[1] = forwardST[0];
		st[2] = forwardST[1];
		st[3] = forwardST[2];
		st[4] = forwardST[3];
		st[5] = forwardST[4];
		st[6] = forwardST[5];
	} else {
		st[1] = backwardST[0];
		st[2] = backwardST[1];
		st[3] = backwardST[2];
		st[4] = backwardST[3];
		st[5] = backwardST[4];
		st[6] = backwardST[5];
	}
}

int main(int argc, char *argv[])
{
	if (ParseArguments(argc, argv) != 0) {
		PrintUsage();
		return 0;
	}

	/* read the configure file */
	ReadConfigureFile(fiber_pathname, &confPathname, &vertPathname, w, h, d, scale, stepSize, maxLength);

	/* compute structure tensor for each voxel */
	float *fsr_st = new float[w*h*d*7];
	memset(fsr_st, 0, sizeof(float)*w*h*d*7);

	CFiber **fibers = new CFiber*[2*(kernel_size+1)+1];

	int offset[9];
    offset[ 0] = 0;		offset[ 1] = -1;	offset[ 2] = 1;
    offset[ 3] = -w-1;	offset[ 4] = -w;	offset[ 5] = -w+1;
    offset[ 6] = w-1;	offset[ 7] = w;		offset[ 8] = w+1;

	float length = fiber_length > maxLength ? maxLength : fiber_length;
	uint id = uint(length / stepSize);
	float delta = 1.0f / (float)scale;

	int ks = 2 * kernel_size + 1;
	float *gaussian = new float[ks*ks*ks];
	for (int z = 0; z < ks; ++z) {
		for (int y = 0; y < ks; ++y) {
			for (int x = 0; x < ks; ++x) {
				int index = (z * ks + y) * ks + x;
				float dis = float((x - kernel_size) * (x - kernel_size) + (y - kernel_size) * (y - kernel_size) + (z - kernel_size) * (z - kernel_size));
				gaussian[index] = exp(-dis / (0.25f * ks * ks));
			}
		}
	}

	omp_set_num_threads(64);

	for (int z = kernel_size + 1; z < d - kernel_size - 1; ++z) {
		int cw, ch;
		for (int k = -kernel_size-1; k <= kernel_size+1; ++k) {
			ReadFibers(confPathname[z-k], vertPathname[z-k], &(fibers[k+kernel_size+1]), cw, ch);
		}
		

		for (int y = kernel_size + 1; y < h - kernel_size - 1; ++y) {
#pragma omp parallel for            
			for (int x = kernel_size + 1; x < w - kernel_size - 1; ++x) {
				int index = y * w + x;
				CFiber *base_fiber = &(fibers[kernel_size+1][index]);

				for (int cz = -kernel_size; cz <= kernel_size; ++cz) {
					for (int cy = -kernel_size; cy <= kernel_size; ++cy) { 
						for (int cx = -kernel_size; cx <= kernel_size; ++cx) {
							int temp_index = (y + cy) * w + (x + cx);
							CFiber *f[27];
							for (int i = 1; i < 9; ++i) {
								f[i] = &(fibers[1+cz+kernel_size][temp_index+offset[i]]);
							}
							for (int i = 0; i < 9; ++i) {
								f[i+9] = &(fibers[0+cz+kernel_size][temp_index+offset[i]]);
							}
							for (int i = 0; i < 9; ++i) {
								f[i+18] = &(fibers[2+cz+kernel_size][temp_index+offset[i]]);
							}

							float temp_tensor[7] = {0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f};
							ComputeFSR_Tensor(f, id, delta, temp_tensor, base_fiber);

							int kx = kernel_size + cx;
							int ky = kernel_size + cy;
							int kz = kernel_size + cz;
							int gk_index = (kz * ks + ky) * ks + kx;

							fsr_st[(z*w*h+index)*7+0] += gaussian[gk_index] * temp_tensor[0];
							fsr_st[(z*w*h+index)*7+1] += gaussian[gk_index] * temp_tensor[1];
							fsr_st[(z*w*h+index)*7+2] += gaussian[gk_index] * temp_tensor[2];
							fsr_st[(z*w*h+index)*7+3] += gaussian[gk_index] * temp_tensor[3];
							fsr_st[(z*w*h+index)*7+4] += gaussian[gk_index] * temp_tensor[4];
							fsr_st[(z*w*h+index)*7+5] += gaussian[gk_index] * temp_tensor[5];
							fsr_st[(z*w*h+index)*7+6] += gaussian[gk_index] * temp_tensor[6];
						}
					}
				}
			}
		}

		for (int k = -kernel_size-1; k <= kernel_size+1; ++k) {
			SafeDeleteArray(fibers[k+kernel_size+1]);
		}

		printf("z = %d\n", z);

	}

	/* save result */
	CDTI *pST = new CDTI();
	pST->CreateDTI(w, h, d);

	for (int z = 0; z < d; ++z) {
		for (int y = 0; y < h; ++y) {
			for (int x = 0; x < w; ++x) {
				int index = (z * h + y) * w + x;
				pST->SetDTI(x, y, z, &(fsr_st[index*7]));
			}
		}
	}
	pST->SaveDTIFile(dst_pathname);
	SafeDelete(pST);

	for (int z = 0; z < d; ++z) {
		SafeDeleteArray(confPathname[z]);
		SafeDeleteArray(vertPathname[z]);
	}
	SafeDeleteArray(confPathname);
	SafeDeleteArray(vertPathname);

	SafeDeleteArray(fsr_st);

	SafeDeleteArray(fibers);
	SafeDeleteArray(gaussian);

	return 0;
}